package db

import (
	"context"
	"errors"
	"time"

	"gitee.com/huangxinbo/xcas/pkg/logger"
	"go.uber.org/zap"
	"go.uber.org/zap/zapcore"
	"gorm.io/gorm"
	gormLogger "gorm.io/gorm/logger"
)

type GormLogger struct {
	Logger                    *logger.Logger
	LogLevel                  gormLogger.LogLevel
	SlowThreshold             time.Duration
	SkipCallerLookup          bool
	IgnoreRecordNotFoundError bool
	Context                   func(ctx context.Context) []zapcore.Field
	EnableSQLLog              bool // 控制是否记录 SQL 语句
}

func NewLogger(logger *logger.Logger, enableSQLLog bool) (*GormLogger, error) {
	return &GormLogger{
		Logger:                    logger,
		LogLevel:                  gormLogger.Info,
		SlowThreshold:             100 * time.Millisecond,
		SkipCallerLookup:          false,
		IgnoreRecordNotFoundError: false,
		Context:                   nil,
		EnableSQLLog:              enableSQLLog,
	}, nil
}

func (l *GormLogger) LogMode(level gormLogger.LogLevel) gormLogger.Interface {
	return &GormLogger{
		Logger:                    l.Logger,
		SlowThreshold:             l.SlowThreshold,
		LogLevel:                  level,
		SkipCallerLookup:          l.SkipCallerLookup,
		IgnoreRecordNotFoundError: l.IgnoreRecordNotFoundError,
		Context:                   l.Context,
		EnableSQLLog:              l.EnableSQLLog,
	}
}

func (l *GormLogger) Info(ctx context.Context, str string, args ...interface{}) {
	if l.LogLevel < gormLogger.Info {
		return
	}
	l.Logger.Info(str, zap.Any("args", args))
}

func (l *GormLogger) Warn(ctx context.Context, str string, args ...interface{}) {
	if l.LogLevel < gormLogger.Warn {
		return
	}
	l.Logger.Warn(str, zap.Any("args", args))
}

func (l *GormLogger) Error(ctx context.Context, str string, args ...interface{}) {
	if l.LogLevel < gormLogger.Error {
		return
	}
	l.Logger.Error(str, zap.Any("args", args))
}

func (l *GormLogger) Trace(ctx context.Context, begin time.Time, fc func() (sql string, rowsAffected int64), err error) {
	if l.LogLevel <= 0 {
		return
	}
	elapsed := time.Since(begin)
	sql, rows := fc()

	// 如果启用了 SQL 日志，每条 SQL 语句都会记录
	if l.EnableSQLLog {
		switch {
		case err != nil && l.LogLevel >= gormLogger.Error && (!l.IgnoreRecordNotFoundError || !errors.Is(err, gorm.ErrRecordNotFound)):
			l.Logger.Error("trace", zap.Error(err), zap.Duration("elapsed", elapsed), zap.Int64("rows", rows), zap.String("sql", sql))
		case l.SlowThreshold != 0 && elapsed > l.SlowThreshold && l.LogLevel >= gormLogger.Warn:
			l.Logger.Warn("trace", zap.Duration("elapsed", elapsed), zap.Int64("rows", rows), zap.String("sql", sql))
		case l.LogLevel >= gormLogger.Info:
			// 正常查询使用 Info 级别记录 SQL 语句，确保在 info 日志级别下也能看到
			l.Logger.Info("trace", zap.Duration("elapsed", elapsed), zap.Int64("rows", rows), zap.String("sql", sql))
		}
	} else {
		// 如果禁用了 SQL 日志，正常和慢查询不记录 SQL，但错误 SQL 仍然要记录
		switch {
		case err != nil && l.LogLevel >= gormLogger.Error && (!l.IgnoreRecordNotFoundError || !errors.Is(err, gorm.ErrRecordNotFound)):
			// 错误 SQL 总是记录 SQL 语句，即使禁用了 SQL 日志
			l.Logger.Error("trace", zap.Error(err), zap.Duration("elapsed", elapsed), zap.Int64("rows", rows), zap.String("sql", sql))
		case l.SlowThreshold != 0 && elapsed > l.SlowThreshold && l.LogLevel >= gormLogger.Warn:
			l.Logger.Warn("trace", zap.Duration("elapsed", elapsed), zap.Int64("rows", rows))
		case l.LogLevel >= gormLogger.Info:
			l.Logger.Debug("trace", zap.Duration("elapsed", elapsed), zap.Int64("rows", rows))
		}
	}
}
